{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f31827b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ langgraph 已安装\n",
      "✓ langchain 已安装\n",
      "✓ pydantic 已安装\n"
     ]
    }
   ],
   "source": [
    "# 安装依赖 (如已安装可跳过)\n",
    "import sys, subprocess, os\n",
    "def ensure(pkg):\n",
    "    try:\n",
    "        __import__(pkg)\n",
    "        print(f'✓ {pkg} 已安装')\n",
    "    except Exception:\n",
    "        print(f'安装 {pkg} ...')\n",
    "        subprocess.check_call([sys.executable, '-m', 'pip', 'install', pkg])\n",
    "for p in ['langgraph', 'langchain', 'pydantic']:\n",
    "    ensure(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19eb8e62",
   "metadata": {},
   "source": [
    "## 1. 定义状态结构与节点函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "15480b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict, List, Optional\n",
    "from langgraph.graph import StateGraph, END\n",
    "import time\n",
    "\n",
    "# 使用 TypedDict 描述状态；也可用 Pydantic BaseModel\n",
    "class FlowState(TypedDict, total=False):\n",
    "    input: str\n",
    "    tokens: List[str]\n",
    "    transformed: List[str]\n",
    "    output: str\n",
    "    trace: List[str]\n",
    "\n",
    "def parse(state: FlowState) -> FlowState:\n",
    "    text = state['input']\n",
    "    tokens = text.replace(',', ' ').split()\n",
    "    trace = state.get('trace', []) + [f'parse:{len(tokens)} tokens']\n",
    "    return {**state, 'tokens': tokens, 'trace': trace}\n",
    "\n",
    "def transform(state: FlowState) -> FlowState:\n",
    "    tokens = state.get('tokens', [])\n",
    "    transformed = [t.upper() for t in tokens]\n",
    "    trace = state.get('trace', []) + ['transform:upper']\n",
    "    return {**state, 'transformed': transformed, 'trace': trace}\n",
    "\n",
    "def finalize(state: FlowState) -> FlowState:\n",
    "    transformed = state.get('transformed', [])\n",
    "    output = ' | '.join(transformed)\n",
    "    trace = state.get('trace', []) + ['finalize']\n",
    "    return {**state, 'output': output, 'trace': trace}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d38a497",
   "metadata": {},
   "source": [
    "## 2. 构建图并编译"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c5d15a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Graph compiled: nodes = ['parse', 'transform', 'finalize']\n"
     ]
    }
   ],
   "source": [
    "graph = StateGraph(FlowState)\n",
    "graph.add_node('parse', parse)\n",
    "graph.add_node('transform', transform)\n",
    "graph.add_node('finalize', finalize)\n",
    "\n",
    "# 边：parse → transform → finalize → END\n",
    "graph.add_edge('parse', 'transform')\n",
    "graph.add_edge('transform', 'finalize')\n",
    "graph.add_edge('finalize', END)\n",
    "\n",
    "graph.set_entry_point('parse')\n",
    "compiled = graph.compile()\n",
    "print('Graph compiled: nodes =', list(graph.nodes.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8cae0d9",
   "metadata": {},
   "source": [
    "## 3. 单次调用示例 (invoke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "55d563ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output: LANGGRAPH | ORCHESTRATION | DEMO | HELLO | WORLD\n",
      "Trace: ['parse:5 tokens', 'transform:upper', 'finalize']\n"
     ]
    }
   ],
   "source": [
    "result = compiled.invoke({'input': 'LangGraph orchestration demo, hello world'})\n",
    "print('Output:', result.get('output'))\n",
    "print('Trace:', result.get('trace'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cfd48e1",
   "metadata": {},
   "source": [
    "## 4. 流式运行 (stream)\n",
    "可逐步查看每个节点对状态的增量更新。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9d5afdd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- event ---\n",
      "{'parse': {'input': 'stream this sentence please', 'tokens': ['stream', 'this', 'sentence', 'please'], 'trace': ['parse:4 tokens']}}\n",
      "--- event ---\n",
      "{'transform': {'input': 'stream this sentence please', 'tokens': ['stream', 'this', 'sentence', 'please'], 'trace': ['parse:4 tokens', 'transform:upper'], 'transformed': ['STREAM', 'THIS', 'SENTENCE', 'PLEASE']}}\n",
      "--- event ---\n",
      "{'finalize': {'input': 'stream this sentence please', 'tokens': ['stream', 'this', 'sentence', 'please'], 'transformed': ['STREAM', 'THIS', 'SENTENCE', 'PLEASE'], 'trace': ['parse:4 tokens', 'transform:upper', 'finalize'], 'output': 'STREAM | THIS | SENTENCE | PLEASE'}}\n"
     ]
    }
   ],
   "source": [
    "for event in compiled.stream({'input': 'stream this sentence please'}):\n",
    "    # event 是一个 (update_dict) 或 (key, value) 取决于版本，这里简单打印\n",
    "    print('--- event ---')\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9759b42c",
   "metadata": {},
   "source": [
    "## 5. 扩展示例：条件分支 (可选)\n",
    "下面演示如何按 token 数量分支：少于 4 走 A，多于等于 4 走 B，然后汇合 finalize。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c85bd91e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Branch graph compiled.\n",
      "{'input': 'few tokens', 'tokens': ['few', 'tokens'], 'transformed': ['A:few', 'A:tokens'], 'output': 'A:few | A:tokens', 'trace': ['parse:2 tokens', 'branch_a', 'finalize']}\n",
      "{'input': 'this sentence has many tokens', 'tokens': ['this', 'sentence', 'has', 'many', 'tokens'], 'transformed': ['B:this', 'B:sentence', 'B:has', 'B:many', 'B:tokens'], 'output': 'B:this | B:sentence | B:has | B:many | B:tokens', 'trace': ['parse:5 tokens', 'branch_b', 'finalize']}\n"
     ]
    }
   ],
   "source": [
    "from typing import Literal\n",
    "\n",
    "def branch_a(state: FlowState) -> FlowState:\n",
    "    trace = state.get('trace', []) + ['branch_a']\n",
    "    return {**state, 'transformed': ['A:' + t for t in state.get('tokens', [])], 'trace': trace}\n",
    "\n",
    "def branch_b(state: FlowState) -> FlowState:\n",
    "    trace = state.get('trace', []) + ['branch_b']\n",
    "    return {**state, 'transformed': ['B:' + t for t in state.get('tokens', [])], 'trace': trace}\n",
    "\n",
    "def router(state: FlowState) -> Literal[\"branch_a\", \"branch_b\"]:\n",
    "    # 返回下一节点名称\n",
    "    if len(state.get('tokens', [])) < 4: return 'branch_a'\n",
    "    return 'branch_b'\n",
    "\n",
    "branch_graph = StateGraph(FlowState)\n",
    "branch_graph.add_node('parse', parse)\n",
    "branch_graph.add_node('branch_a', branch_a)\n",
    "branch_graph.add_node('branch_b', branch_b)\n",
    "branch_graph.add_node('finalize', finalize)\n",
    "branch_graph.add_conditional_edges('parse', router,{\n",
    "    \"branch_a\": \"branch_a\",\n",
    "    \"branch_b\": \"branch_b\"\n",
    "})  # 可选：显式映射返回值到节点（若返回值就是节点名，可省略此字典）\n",
    "branch_graph.add_edge('branch_a', 'finalize')\n",
    "branch_graph.add_edge('branch_b', 'finalize')\n",
    "branch_graph.add_edge('finalize', END)\n",
    "branch_graph.set_entry_point('parse')\n",
    "compiled_branch = branch_graph.compile()\n",
    "print('Branch graph compiled.')\n",
    "print(compiled_branch.invoke({'input': 'few tokens'}))\n",
    "print(compiled_branch.invoke({'input': 'this sentence has many tokens'}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a55df8",
   "metadata": {},
   "source": [
    "1. 使用 StateGraph 定义节点与边。\n",
    "2. 使用 `invoke` 获得最终状态。\n",
    "3. 使用 `stream` 观察中间状态更新。\n",
    "4. 动态路由 (条件分支) 的简单写法。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
